package com.demo.LSTM;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.api.Layer;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.*;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.LSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.evaluation.regression.RegressionEvaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.ArrayList;
import java.util.List;

public class LSTMLoadPrediction1Demo {

    public static void main(String[] args) {
        // 负载和天气数据
        int numInputs = 2;
        int numOutputs = 1;
        int numHiddenUnits = 10;
        int numEpochs = 100;

        // 创建数据集
        //double[] inputArray = {0.1, 0.2, 0.3, 0.4};
        //double[] outputArray = {0.2, 0.3, 0.4, 0.5};
        // 创建数据集
        double[] loadArray  = {0.1, 0.2, 0.3, 0.4};
        double[] temperatureArray  = {20.0, 21.0, 22.0, 23.0};

        List<DataSet> dataList = new ArrayList<>(8);
        //for (int i = 0; i < inputArray.length; i++) {
        for (int i = 0; i < loadArray.length; i++) {
            //double[] input = {inputArray[i]};
            //double[] output = {outputArray[i]};
            //dataList.add(new DataSet(Nd4j.create(input, new int[]{1, numInputs}), Nd4j.create(output, new int[]{1, numOutputs})));
            //dataList.add(new DataSet(Nd4j.create(input, new int[]{1, numInputs, 1}), Nd4j.create(output, new int[]{1, numOutputs, 1})));

            double[] input = {loadArray[i], temperatureArray[i]};
            double[] output = {loadArray[i]};
            dataList.add(new DataSet(Nd4j.create(input, new int[]{1, numInputs, 1}), Nd4j.create(output, new int[]{1, numOutputs, 1})));
        }

        ListDataSetIterator<DataSet> iterator = new ListDataSetIterator<>(dataList, 1);

        // 配置神经网络
        MultiLayerConfiguration configuration = new NeuralNetConfiguration.Builder()
                .seed(123)
                .weightInit(WeightInit.XAVIER)
                //.updater(new Adam(0.01))
                .updater(new Adam())
                .list()
                .layer(new LSTM.Builder().nIn(numInputs).nOut(numHiddenUnits).activation(Activation.TANH).build())
                .layer(new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE).activation(Activation.IDENTITY).nIn(numHiddenUnits).nOut(numOutputs).build())

                //.layer(new LSTM.Builder().nIn(numInputs).nOut(numHiddenUnits).activation(Activation.TANH).build())
                //.layer(new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE).activation(Activation.IDENTITY).nIn(numHiddenUnits).nOut(numOutputs).build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(configuration);
        model.init();

        // 训练模型
        for (int i = 0; i < numEpochs; i++) {
            model.fit(iterator);
            iterator.reset();
        }

        // 预测
        //double[] input = {0.5, 0.6, 0.7, 0.2};
        //INDArray inputArr = Nd4j.create(input, new int[]{4, numInputs});
        //INDArray outputArr = model.rnnTimeStep(inputArr);
        //System.out.println("预测结果: " + outputArr);
        // 进行预测
        double[] inputArray = {0.5, 24.0}; // 当前负载和天气数据
        INDArray input = Nd4j.create(inputArray, new int[]{1, numInputs, 1});
        INDArray predictedOutput = model.rnnTimeStep(input);
        System.out.println("Predicted load: " + predictedOutput.getDouble(0));

    }
}
